package com.bigdata.spark.listeners;

import java.util.HashMap;
import java.util.Map;
import org.apache.spark.scheduler.SparkListener;
import org.apache.spark.scheduler.SparkListenerEvent;
import org.apache.spark.sql.SparkSession;
import org.apache.spark.sql.execution.ui.SparkListenerSQLExecutionEnd;
import org.apache.spark.sql.execution.ui.SparkListenerSQLExecutionStart;
import scala.Option;

/**
 * DESC 通过spark 监听器实现采集sql审计
 * <p>
 * spark.extraListeners=com.bigdata.spark.listeners.CustomSparkSQLListener
 * </p>
 *
 * @author Cyber
 * <p> Created By 2025/5/22
 * @version 1.0
 */

public class CustomSparkSQLListener extends SparkListener {
    // 存储 SQL 执行元数据（ExecutionID -> SQL）
    private final Map<Long, String> sqlExecutionMap = new HashMap<>();

    @Override
    public void onOtherEvent(SparkListenerEvent event) {
        // 监听所有事件，筛选出 SQL 相关的事件
        if (event instanceof SparkListenerSQLExecutionStart) {
            handleSQLStart((SparkListenerSQLExecutionStart) event);
        } else if (event instanceof SparkListenerSQLExecutionEnd) {
            handleSQLEnd((SparkListenerSQLExecutionEnd) event);
        }
    }

    // 处理 SQL 开始执行事件
    private void handleSQLStart(SparkListenerSQLExecutionStart startEvent) {
        long executionId = startEvent.executionId();
        String sql = "";  // 获取 SQL
        String user = SparkSession.active().sparkContext().sparkUser();
        try {
            // 使用反射获取私有字段
            java.lang.reflect.Field field = startEvent.getClass().getDeclaredField("description");
            field.setAccessible(true);
            sql = (String) field.get(startEvent);
        } catch (Exception e) {
            e.printStackTrace();
        }
        sqlExecutionMap.put(executionId, sql);

        System.out.printf("[SQL Start] Execution ID: %d, User: %s, SQL: %s %n",
                executionId, user, sql);
    }

    // 处理 SQL 执行结束事件
    private void handleSQLEnd(SparkListenerSQLExecutionEnd endEvent) {
        long executionId = endEvent.executionId();
        endEvent.time();
        Long executionDuration = Option.apply(endEvent.duration()).getOrElse(() -> -1L);
        String status = endEvent.executionFailure().isDefined() ? "FAILED" : "SUCCESS";
        String sql = sqlExecutionMap.getOrDefault(executionId, "UNKNOWN");
        // 清理元数据
        sqlExecutionMap.remove(executionId);

        // 输出日志或写入外部系统
        System.out.printf("[SQL End] Execution ID: %d, SQL: %s, Status: %s, executionDuration:%d %n",
                executionId, sql, status, executionDuration);
    }
}
